import torch
import torch.nn as nn
from model.utils import ConvBlock, DownSampling, UpSampling
from dataloader.dataload import getDataloader
from model.csp_module import C3, Conv, dark_module


class Csp_unet_model(nn.Module):
    def __init__(self, ConvBlock, C3, Conv, UpSampling, base_channels=64, base_depth=1):
        super(Csp_unet_model, self).__init__()
        self.conv_block1 = ConvBlock(in_ch=3, out_ch=64, kernel_size=(3, 3), stride=1, padding=1)
        self.down1 = nn.Sequential(
            # 320, 320, 64 -> 160, 160, 128
            Conv(base_channels, base_channels * 2, 3, 2),
            # 160, 160, 128 -> 160, 160, 128
            C3(base_channels * 2, base_channels * 2, base_depth),
        )
        self.down2 = nn.Sequential(
            Conv(base_channels * 2, base_channels * 4, 3, 2),
            C3(base_channels * 4, base_channels * 4, base_depth * 3),
        )
        self.down3 = nn.Sequential(
            Conv(base_channels * 4, base_channels * 8, 3, 2),
            C3(base_channels * 8, base_channels * 8, base_depth * 3),
        )
        self.conv_block4 = ConvBlock(in_ch=512, out_ch=1024, kernel_size=(3, 3), stride=2, padding=1)
        self.up4 = UpSampling(1024, 512, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h4 = ConvBlock(in_ch=1024, out_ch=512, kernel_size=(3, 3), stride=1, padding=1)
        self.up3 = UpSampling(512, 256, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h3 = ConvBlock(in_ch=512, out_ch=256, kernel_size=(3, 3), stride=1, padding=1)
        self.up2 = UpSampling(256, 128, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h2 = ConvBlock(in_ch=256, out_ch=128, kernel_size=(3, 3), stride=1, padding=1)
        self.up1 = UpSampling(128, 64, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h1 = ConvBlock(in_ch=128, out_ch=64, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h0 = nn.Conv2d(64, 3, (3, 3), 1, 1)

    def forward(self, x):
        output1 = self.conv_block1(x)  # 416x416x64
        output2 = self.down1(output1)  # 208x208x128
        output3 = self.down2(output2)  # 104x104x256
        output4 = self.down3(output3)  # 52x52x512
        output5 = self.conv_block4(output4) # 26x26x1024
        up4 = self.conv_h4(self.up4(output5, output4))# 52x52x512
        up3 = self.conv_h3(self.up3(up4, output3))# 104x104x256
        up2 = self.conv_h2(self.up2(up3, output2))# 208x208x128
        up1 = self.conv_h1(self.up1(up2, output1))# 416x416x64
        up0 = self.conv_h0(up1)
        return up0


class Unet_model(nn.Module):
    def __init__(self, ConvBlock, DownSampling, UpSampling):
        super(Unet_model, self).__init__()
        self.conv_block1 = ConvBlock(in_ch=3, out_ch=64, kernel_size=(3, 3), stride=1, padding=1)
        self.down1 = DownSampling(kernel_size=(2, 2), stride=2)
        self.conv_block2 = ConvBlock(in_ch=64, out_ch=128, kernel_size=(3, 3), stride=1, padding=1)
        self.down2 = DownSampling(kernel_size=(2, 2), stride=2)
        self.conv_block3 = ConvBlock(in_ch=128, out_ch=256, kernel_size=(3, 3), stride=1, padding=1)
        self.down3 = DownSampling(kernel_size=(2, 2), stride=2)
        self.conv_block4 = ConvBlock(in_ch=256, out_ch=512, kernel_size=(3, 3), stride=1, padding=1)
        self.down4 = DownSampling(kernel_size=(2, 2), stride=2)
        self.conv_block5 = ConvBlock(in_ch=512, out_ch=1024, kernel_size=(3, 3), stride=1, padding=1)
        self.up4 = UpSampling(1024, 512, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h4 = ConvBlock(in_ch=1024, out_ch=512, kernel_size=(3, 3), stride=1, padding=1)
        self.up3 = UpSampling(512, 256, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h3 = ConvBlock(in_ch=512, out_ch=256, kernel_size=(3, 3), stride=1, padding=1)
        self.up2 = UpSampling(256, 128, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h2 = ConvBlock(in_ch=256, out_ch=128, kernel_size=(3, 3), stride=1, padding=1)
        self.up1 = UpSampling(128, 64, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h1 = ConvBlock(in_ch=128, out_ch=64, kernel_size=(3, 3), stride=1, padding=1)
        self.conv_h0 = nn.Conv2d(64, 3, (3, 3), 1, 1)

    def forward(self, x):
        output1 = self.conv_block1(x)
        output2 = self.conv_block2(self.down1(output1))
        output3 = self.conv_block3(self.down2(output2))
        output4 = self.conv_block4(self.down3(output3))
        output5 = self.conv_block5(self.down4(output4))
        up4 = self.conv_h4(self.up4(output5, output4))
        up3 = self.conv_h3(self.up3(up4, output3))
        up2 = self.conv_h2(self.up2(up3, output2))
        up1 = self.conv_h1(self.up1(up2, output1))
        up0 = self.conv_h0(up1)
        return up0


if __name__ == '__main__':
    path = "../data/data"
    train_loader, val_loader, test_loader = getDataloader(path, BatchSize=2, split=[0.1, 0.8, 0.1])
    device = "cuda" if torch.cuda.is_available() else "cpu"
    img, mask = iter(train_loader).next()[0].to(device), iter(train_loader).next()[1].to(device)
    # model = Unet_model(ConvBlock, DownSampling, UpSampling).to(device)
    model = Csp_unet_model(ConvBlock, C3, Conv, UpSampling, base_channels=64, base_depth=1)
    output = model(img)
    print(output.shape)
    print(mask.shape)
    print(output)
